"""
11_centrality_algorithms.py
===========================
Centrality Algorithms with GDS

This module covers:
1. PageRank algorithm
2. Betweenness Centrality
3. Closeness Centrality
4. Degree Centrality
5. Eigenvector Centrality
6. Article Rank

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from neo4j import Driver
import pandas as pd
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


@dataclass
class CentralityResult:
    """Result from centrality algorithm."""
    node_id: str
    node_name: str
    score: float
    rank: int


class CentralityAlgorithms:
    """
    Implements various centrality algorithms using Neo4j GDS.
    
    Centrality algorithms help identify the most important nodes in a graph.
    """
    
    def __init__(self, driver: Driver):
        self.driver = driver
        self.graph_name = "centrality_graph"
    
    def create_graph_projection(
        self,
        node_labels: List[str] = None,
        relationship_types: List[str] = None
    ) -> Dict[str, Any]:
        """
        Create a graph projection for centrality analysis.
        
        Args:
            node_labels: Node labels to include
            relationship_types: Relationship types to include
            
        Returns:
            Projection statistics
        """
        node_labels = node_labels or ["Patient", "Disease", "Medication"]
        relationship_types = relationship_types or [
            "DIAGNOSED_WITH", "TAKES", "TREATS", "INTERACTS_WITH"
        ]
        
        # Drop existing projection
        drop_query = f"""
        CALL gds.graph.drop('{self.graph_name}', false)
        YIELD graphName
        """
        
        # Create new projection
        create_query = """
        CALL gds.graph.project(
            $graph_name,
            $node_labels,
            {
                DIAGNOSED_WITH: {orientation: 'UNDIRECTED'},
                TAKES: {orientation: 'UNDIRECTED'},
                TREATS: {orientation: 'UNDIRECTED'},
                INTERACTS_WITH: {orientation: 'UNDIRECTED'}
            }
        )
        YIELD graphName, nodeCount, relationshipCount
        RETURN graphName, nodeCount, relationshipCount
        """
        
        with self.driver.session() as session:
            try:
                session.run(drop_query)
            except:
                pass
            
            try:
                result = session.run(
                    create_query,
                    graph_name=self.graph_name,
                    node_labels=node_labels
                )
                record = result.single()
                return record.data() if record else None
            except Exception as e:
                print(f"   Projection failed: {e}")
                return None
    
    # =========================================================
    # PAGERANK
    # =========================================================
    
    def pagerank(
        self,
        damping_factor: float = 0.85,
        max_iterations: int = 20,
        tolerance: float = 0.0001,
        write_property: str = None
    ) -> List[Dict[str, Any]]:
        """
        Run PageRank algorithm.
        
        PageRank measures the importance of nodes based on the structure
        of incoming links. Originally developed by Google for web pages.
        
        Args:
            damping_factor: Probability of following a link (0.85 typical)
            max_iterations: Maximum iterations
            tolerance: Convergence tolerance
            write_property: If provided, write results to this property
            
        Returns:
            List of nodes with PageRank scores
            
        Use Cases:
            - Find most influential patients (many connections)
            - Identify key diseases (connected to many treatments)
            - Discover important medications
        """
        if write_property:
            query = """
            CALL gds.pageRank.write($graph_name, {
                dampingFactor: $damping,
                maxIterations: $max_iter,
                tolerance: $tol,
                writeProperty: $write_prop
            })
            YIELD nodePropertiesWritten, ranIterations, didConverge
            RETURN nodePropertiesWritten, ranIterations, didConverge
            """
            
            with self.driver.session() as session:
                result = session.run(
                    query,
                    graph_name=self.graph_name,
                    damping=damping_factor,
                    max_iter=max_iterations,
                    tol=tolerance,
                    write_prop=write_property
                )
                return result.single().data()
        else:
            query = """
            CALL gds.pageRank.stream($graph_name, {
                dampingFactor: $damping,
                maxIterations: $max_iter,
                tolerance: $tol
            })
            YIELD nodeId, score
            WITH gds.util.asNode(nodeId) AS node, score
            RETURN 
                node.id AS id,
                node.name AS name,
                labels(node)[0] AS label,
                score
            ORDER BY score DESC
            LIMIT 20
            """
            
            with self.driver.session() as session:
                result = session.run(
                    query,
                    graph_name=self.graph_name,
                    damping=damping_factor,
                    max_iter=max_iterations,
                    tol=tolerance
                )
                return [record.data() for record in result]
    
    # =========================================================
    # BETWEENNESS CENTRALITY
    # =========================================================
    
    def betweenness_centrality(
        self,
        sampling_size: int = None,
        sampling_seed: int = None
    ) -> List[Dict[str, Any]]:
        """
        Run Betweenness Centrality algorithm.
        
        Betweenness measures how often a node lies on shortest paths
        between other nodes. High betweenness = bridge/broker node.
        
        Args:
            sampling_size: Sample size for approximation (faster)
            sampling_seed: Random seed for reproducibility
            
        Returns:
            List of nodes with betweenness scores
            
        Use Cases:
            - Find medications that bridge multiple conditions
            - Identify patients connecting different disease clusters
            - Discover key intermediary nodes
        """
        config = {}
        if sampling_size:
            config['samplingSize'] = sampling_size
        if sampling_seed:
            config['samplingSeed'] = sampling_seed
        
        query = """
        CALL gds.betweenness.stream($graph_name, $config)
        YIELD nodeId, score
        WITH gds.util.asNode(nodeId) AS node, score
        WHERE score > 0
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            score
        ORDER BY score DESC
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                config=config
            )
            return [record.data() for record in result]
    
    # =========================================================
    # CLOSENESS CENTRALITY
    # =========================================================
    
    def closeness_centrality(
        self,
        use_wasserman_faust: bool = False
    ) -> List[Dict[str, Any]]:
        """
        Run Closeness Centrality algorithm.
        
        Closeness measures how close a node is to all other nodes.
        High closeness = can reach other nodes quickly.
        
        Args:
            use_wasserman_faust: Use Wasserman-Faust formula for
                                 disconnected graphs
            
        Returns:
            List of nodes with closeness scores
            
        Use Cases:
            - Find diseases that are "close" to many treatments
            - Identify well-connected patients
            - Discover central hub medications
        """
        query = """
        CALL gds.closeness.stream($graph_name, {
            useWassermanFaust: $wf
        })
        YIELD nodeId, score
        WITH gds.util.asNode(nodeId) AS node, score
        WHERE score > 0
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            score
        ORDER BY score DESC
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                wf=use_wasserman_faust
            )
            return [record.data() for record in result]
    
    # =========================================================
    # DEGREE CENTRALITY
    # =========================================================
    
    def degree_centrality(
        self,
        orientation: str = "UNDIRECTED"
    ) -> List[Dict[str, Any]]:
        """
        Run Degree Centrality algorithm.
        
        Degree centrality simply counts connections. Fast and intuitive.
        
        Args:
            orientation: NATURAL, REVERSE, or UNDIRECTED
            
        Returns:
            List of nodes with degree scores
            
        Use Cases:
            - Find patients with most diagnoses
            - Identify medications treating most conditions
            - Discover diseases with most treatment options
        """
        query = """
        CALL gds.degree.stream($graph_name, {
            orientation: $orientation
        })
        YIELD nodeId, score
        WITH gds.util.asNode(nodeId) AS node, score
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            toInteger(score) AS degree
        ORDER BY score DESC
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                orientation=orientation
            )
            return [record.data() for record in result]
    
    # =========================================================
    # EIGENVECTOR CENTRALITY
    # =========================================================
    
    def eigenvector_centrality(
        self,
        max_iterations: int = 20,
        tolerance: float = 0.0001
    ) -> List[Dict[str, Any]]:
        """
        Run Eigenvector Centrality algorithm.
        
        Similar to PageRank but without damping. Measures influence
        based on connections to other influential nodes.
        
        Args:
            max_iterations: Maximum iterations
            tolerance: Convergence tolerance
            
        Returns:
            List of nodes with eigenvector scores
        """
        query = """
        CALL gds.eigenvector.stream($graph_name, {
            maxIterations: $max_iter,
            tolerance: $tol
        })
        YIELD nodeId, score
        WITH gds.util.asNode(nodeId) AS node, score
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            score
        ORDER BY score DESC
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                max_iter=max_iterations,
                tol=tolerance
            )
            return [record.data() for record in result]
    
    # =========================================================
    # ARTICLE RANK
    # =========================================================
    
    def article_rank(
        self,
        damping_factor: float = 0.85,
        max_iterations: int = 20
    ) -> List[Dict[str, Any]]:
        """
        Run Article Rank algorithm.
        
        Variant of PageRank that lowers influence of low-degree nodes.
        Better for certain graph types.
        
        Args:
            damping_factor: Damping factor
            max_iterations: Maximum iterations
            
        Returns:
            List of nodes with ArticleRank scores
        """
        query = """
        CALL gds.articleRank.stream($graph_name, {
            dampingFactor: $damping,
            maxIterations: $max_iter
        })
        YIELD nodeId, score
        WITH gds.util.asNode(nodeId) AS node, score
        RETURN 
            node.id AS id,
            node.name AS name,
            labels(node)[0] AS label,
            score
        ORDER BY score DESC
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                damping=damping_factor,
                max_iter=max_iterations
            )
            return [record.data() for record in result]
    
    # =========================================================
    # COMPARISON
    # =========================================================
    
    def compare_centralities(self) -> pd.DataFrame:
        """
        Compare all centrality algorithms on the same graph.
        
        Returns:
            DataFrame with centrality scores for comparison
        """
        results = {}
        
        # Get PageRank
        pr_results = self.pagerank()
        for r in pr_results:
            key = r['id'] or r['name']
            if key not in results:
                results[key] = {'name': r['name'], 'label': r['label']}
            results[key]['pagerank'] = r['score']
        
        # Get Degree
        deg_results = self.degree_centrality()
        for r in deg_results:
            key = r['id'] or r['name']
            if key not in results:
                results[key] = {'name': r['name'], 'label': r['label']}
            results[key]['degree'] = r['degree']
        
        # Get Betweenness
        bet_results = self.betweenness_centrality()
        for r in bet_results:
            key = r['id'] or r['name']
            if key not in results:
                results[key] = {'name': r['name'], 'label': r['label']}
            results[key]['betweenness'] = r['score']
        
        # Convert to DataFrame
        df = pd.DataFrame.from_dict(results, orient='index')
        df = df.fillna(0)
        
        # Normalize scores for comparison
        for col in ['pagerank', 'degree', 'betweenness']:
            if col in df.columns:
                max_val = df[col].max()
                if max_val > 0:
                    df[f'{col}_normalized'] = df[col] / max_val
        
        return df.sort_values('pagerank', ascending=False)
    
    def cleanup(self):
        """Drop the graph projection."""
        query = f"CALL gds.graph.drop('{self.graph_name}', false)"
        with self.driver.session() as session:
            try:
                session.run(query)
            except:
                pass


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Centrality Algorithms - Demonstration")
    print("=" * 60)
    
    driver = get_driver()
    centrality = CentralityAlgorithms(driver)
    
    # Create projection
    print("\n📌 Creating Graph Projection")
    print("-" * 40)
    proj = centrality.create_graph_projection()
    if proj:
        print(f"   Nodes: {proj['nodeCount']}")
        print(f"   Relationships: {proj['relationshipCount']}")
    else:
        print("   Using fallback - limited data")
    
    # PageRank
    print("\n📌 PageRank Results")
    print("-" * 40)
    try:
        pr_results = centrality.pagerank()
        for i, r in enumerate(pr_results[:5], 1):
            print(f"   {i}. {r['name']} ({r['label']}): {r['score']:.6f}")
    except Exception as e:
        print(f"   Error: {e}")
    
    # Betweenness
    print("\n📌 Betweenness Centrality Results")
    print("-" * 40)
    try:
        bet_results = centrality.betweenness_centrality()
        for i, r in enumerate(bet_results[:5], 1):
            print(f"   {i}. {r['name']} ({r['label']}): {r['score']:.6f}")
    except Exception as e:
        print(f"   Error: {e}")
    
    # Degree
    print("\n📌 Degree Centrality Results")
    print("-" * 40)
    try:
        deg_results = centrality.degree_centrality()
        for i, r in enumerate(deg_results[:5], 1):
            print(f"   {i}. {r['name']} ({r['label']}): {r['degree']} connections")
    except Exception as e:
        print(f"   Error: {e}")
    
    # Cleanup
    centrality.cleanup()
    driver.close()
    
    print("\n" + "=" * 60)
    print("✅ Centrality algorithms demonstration complete!")
    print("=" * 60)
